import os

from pydub import AudioSegment, silence
from tqdm import tqdm


def split_audio(
    input_audio_path,
    output_directory,
    max_chunk_duration=10_000,
    min_silence_duration=900,
):
    # Load the audio file
    audio = AudioSegment.from_file(input_audio_path)

    # Perform voice activity detection using pydub's silence module
    segments = silence.detect_silence(
        audio, min_silence_len=min_silence_duration, silence_thresh=-50
    )

    # Split the audio into chunks based on silence segments, avoiding chunks longer than max_chunk_duration
    chunk_number = 1
    start_time = 0
    for start, end in segments:
        if (end - start) <= max_chunk_duration:
            chunk = audio[start_time:end]
            chunk_duration = len(chunk)

            # Check if adding the current segment would exceed the max_chunk_duration
            if chunk_duration + (end - start_time) > max_chunk_duration:
                chunk = audio[start_time : start_time + max_chunk_duration]

            # Save the chunk to WAV format
            chunk_wav = os.path.join(output_directory, f"chunk_{chunk_number}.wav")
            if not os.path.exists(chunk_wav):
                chunk.export(chunk_wav, format="wav")
            else:
                print(f"Chunk {chunk_number} already exists")

            print(
                f"Chunk {chunk_number}: {start_time} ms - {start_time + len(chunk)} ms"
            )
            chunk_number += 1
            start_time += len(chunk)


def split_audio_directory(
    input_directory,
    output_directory,
    max_chunk_duration=10_000,
    min_silence_duration=900,
    audio_extensions=(".wav", ".mp3", ".flac", ".ogg", ".m4a"),
):
    """Process all audio files in a directory"""
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    for filename in tqdm(os.listdir(input_directory)):
        if filename.lower().endswith(audio_extensions):
            input_path = os.path.join(input_directory, filename)
            file_output_dir = os.path.join(
                output_directory, os.path.splitext(filename)[0]
            )

            if not os.path.exists(file_output_dir):
                os.makedirs(file_output_dir)

            split_audio(
                input_path, file_output_dir, max_chunk_duration, min_silence_duration
            )